#pragma once

#include <cstdlib>
#include <cstring>
#include <algorithm>
#include <stdlib.h>

#include "impl.hh"
#include "../meta/constraints.hh"

namespace icesp
{
    namespace heap_impl
    {
        struct alignas(64) heap_info
        {
            impl::item* heap;
            int len;
        };
    }

    struct heaps : meta::non_transferable
    {
        heaps(int n, int nthread, int const* node_count)
            : nthread(nthread)
        {
            // FIXME posix malloc aligned alloc
            // if (posix_memalign((void**)&index, 1024 * 1024, sizeof(int) * n)) {
            //     // TODO err
            // } else {
            // }

            // FIXME malloc aligned alloc
            // index = (int*)_mm_malloc(sizeof(int) * n, 1024 * 1024);
            // hi = (heap_impl::heap_info*)_mm_malloc(sizeof(heap_impl::heap_info) * nthread, 64);
            index = new int[n];
            // hi = new heap_impl::heap_info[nthread];

            // FIXME posix malloc aligned alloc
            auto error = posix_memalign((void**)&hi, 64, sizeof(heap_impl::heap_info) * nthread);
            // TODO handle with error

            std::memset(index, 0xff, sizeof(int) * n);
            for (int i = 0; i < nthread; i++) {
                // FIXME malloc aligned alloc
                //hi[i].heap = (impl::item*)_mm_malloc(sizeof(impl::item) * node_count[i], 64);
                //hi[i].len = 0;
                // hi[i].heap = new impl::item[node_count[i]];
                // FIXME posix malloc aligned alloc
                auto error = posix_memalign((void**)&hi[i].heap, 64, sizeof(impl::item) * node_count[i]);
                hi[i].len = 0;
            }
        }
        ~heaps()
        {
            // FIXME aligned alloc
            //for (int i = 0; i < nthread; i++)
            //    _mm_free(hi[i].heap);
            //_mm_free(index);
            //_mm_free(hi);
            delete[] index;

            // delete[] hi;
            // FIXME posix_memalign

            for (int i = 0; i < nthread; i++) {
                // delete[] hi[i].heap;
                // FIXME posix_memalign
                if (hi[i].heap != NULL) free(hi[i].heap);
            }

            free(hi);
        }

        void add(int thread_id, int key, int value)
        {
            impl::item* & data = hi[thread_id].heap;
            int& len = hi[thread_id].len;
            data[len].key = key;
            data[len].value = value;
            index[key] = len;
            len++;
            up(hi[thread_id], len - 1);
        }

        void update(int thread_id, int key, int value)
        {
            impl::item* & data = hi[thread_id].heap;

            data[index[key]].value = value;
            up(hi[thread_id], index[key]);
        }

        void add_or_update(int thread_id, int key, int value)
        {
            impl::item* & data = hi[thread_id].heap;
            if (index[key] == -1)
                add(thread_id, key, value);
            else if (value < data[index[key]].value)
                update(thread_id, key, value);
        }

        void remove_top(int thread_id, impl::item* d)
        {
            impl::item* & data = hi[thread_id].heap;
            int &len = hi[thread_id].len;

            len--;
            /* hi[thread_id].nswap++; */
            std::swap(data[0], data[len]);
            if (d != nullptr) *d = data[len];
            index[data[0].key] = 0;
            index[data[len].key] = -1;
            down(hi[thread_id], 0);
        }

        impl::item const& top(int thread_id) const
        {
            return hi[thread_id].heap[0];
        }

        bool empty(int thread_id) const
        {
            return hi[thread_id].len == 0;
        }

        int size(int thread_id) const
        {
            return hi[thread_id].len;
        }

    private:
        int father(int x) { return (x - 1) / 2; }

        void up(heap_impl::heap_info& hi, int x)
        {
            auto& data = hi.heap;
            while (x > 0) {
                auto tmp = (x - 1) / 2;
                if (data[tmp].value > data[x].value) {
                    std::swap(data[tmp], data[x]);
                    index[data[tmp].key] = tmp;
                    index[data[x].key] = x;
                    x = tmp;
                } else
                    break;
            }
        }

        void down(heap_impl::heap_info& hi, int x)
        {
            auto& data = hi.heap;
            auto& len = hi.len;
            while (x * 2 + 1 < len) {
                auto tmp = (x * 2 + 2 >= len || data[x * 2 + 1].value < data[x * 2 + 2].value) ?
                    (x * 2 + 1) : x * 2 + 2;
                if (data[tmp].value < data[x].value) {
                    std::swap(data[tmp], data[x]);
                    index[data[tmp].key] = tmp;
                    index[data[x].key] = x;
                    x = tmp;
                } else
                    break;
            }
        }

        int nthread;
        int* index;
        heap_impl::heap_info* hi;
    };


    struct alignas(64) node
    {
        int w[4];
        int to[4];
        int to_color[4];
        int dist;
        int heap_index;
        bool frontier;
    };

    struct lean_heap
    {
        lean_heap(int cap, node* nodes) : cap(cap), size(0), nodes(nodes)
        {
            auto error = posix_memalign((void**)&heap, 64, sizeof(impl::item) * cap);
            std::memset(heap, 0, sizeof(impl::item) * cap);
        }

        ~lean_heap()
        {
            free(heap);
        }

        void up(int x)
        {
            while (x > 0) {
                auto tmp = (x - 1) / 2;
                if (heap[tmp].value > heap[x].value) {
                    std::swap(heap[tmp], heap[x]);
                    nodes[heap[x].key].heap_index = x;
                    x = tmp;
                } else
                    break;
            }
            nodes[heap[x].key].heap_index = x;
        }

        void down(int x)
        {
            auto xval = heap[x].value;
            auto xkey = heap[x].key;
            while (2 * x + 1 < size) {
                auto tmp = (x * 2 + 2 >= size ||
                            heap[x * 2 + 1].value < heap[x * 2 + 2].value)
                            ? x * 2 + 1 : x * 2 + 2;
                if (heap[tmp].value < xval) {
                    heap[x] = heap[tmp];
                    nodes[heap[x].key].heap_index = x;
                    x = tmp;
                } else
                    break;
            }
            heap[x].value = xval; heap[x].key = xkey;
            nodes[heap[x].key].heap_index = x;
        }

        void add(int key, int value)
        {
            if (size == cap) {
                realloc(cap * 2);
                cap *= 2;
            }
            heap[size] = {key, value};
            up(size++);
        }

        void update(int key, int value)
        {
            heap[nodes[key].heap_index].value = value;
            up(nodes[key].heap_index);
        }

        void add_or_update(int key, int value)
        {
            if (nodes[key].heap_index == -1)
                add(key, value);
            else if (value < heap[nodes[key].heap_index].value)
                update(key, value);
        }

        void remove_top(impl::item* pdata)
        {
            size--;
            std::swap(heap[0], heap[size]);
            nodes[heap[size].key].heap_index = -1;
            if (size) down(0);
        }

        void realloc(int new_size)
        {
            impl::item* tmp{nullptr};
            auto error = posix_memalign((void**)&tmp, 64, sizeof(impl::item) * new_size);
            std::copy(heap, heap + cap, tmp);
            free(heap);
            heap = tmp;
        }

        int cap;
        int size;
        node* nodes;
        impl::item* heap;
    };
}

